package com.googlecode.protoserv;

import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.net.Socket;
import java.net.URL;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.zip.DeflaterInputStream;
import java.util.zip.DeflaterOutputStream;
import java.util.zip.GZIPInputStream;
import java.util.zip.GZIPOutputStream;

import com.google.protobuf.ByteString;
import com.google.protobuf.CodedInputStream;
import com.google.protobuf.ExtensionRegistry;
import com.google.protobuf.WireFormat;
import com.googlecode.protoserv.Protocol.RequestProto;
import com.googlecode.protoserv.Protocol.ResponseProto;
import com.googlecode.protoserv.Protocol.RpcStatus;
import com.googlecode.protoserv.Transport.Encoding;
import com.googlecode.protoserv.Transport.SocketProtocolVersion;
import com.googlecode.protoserv.Transport.SocketRequestProto;
import com.googlecode.protoserv.Transport.SocketRequestProto.EncodedRequestProto;
import com.googlecode.protoserv.Transport.SocketResponseProto;
import com.googlecode.protoserv.Transport.SocketResponseProto.EncodedResponseProto;
import com.googlecode.protoserv.Transport.SocketStreamProto;

@SuppressWarnings("rawtypes")
class TcpTransport extends RpcTransport implements Runnable {

  public static final List<Encoding> ACCEPTED_ENCODINGS = Arrays.asList(
      Encoding.DEFLATE, Encoding.GZIP);

  private Object lock;
  private long callId;
  private HashMap<Long, Rpc> calls;
  private ExtensionRegistry registry;
  private URL url;
  private boolean end;

  private Socket socket;
  private InputStream input;
  private OutputStream output;

  public TcpTransport(URL url) {
    this.lock = new Object();
    this.calls = new HashMap<Long, Rpc>();
    this.registry = ExtensionRegistry.newInstance();
    this.url = url;
  }

  @Override
  public void open() throws IOException {
    synchronized (this.lock) {
      if (this.socket != null) {
        throw new IOException("Connection already opened.");
      }
      try {
        this.socket = new Socket(this.url.getHost(), this.url.getPort());
        this.input = this.socket.getInputStream();
        this.output = this.socket.getOutputStream();
      } catch (IOException e) {
        close();
        throw e;
      }
    }
  }

  @Override
  public boolean isOpen() {
    synchronized (this.lock) {
      return this.socket != null;
    }
  }

  private static EncodedRequestProto pack(Encoding encoding,
      RequestProto request) throws IOException {
    EncodedRequestProto.Builder encodedMessage = EncodedRequestProto
        .newBuilder();
    switch (encoding.getNumber()) {
      case Encoding.IDENTITY_VALUE:
        encodedMessage.setIdentity(request);
        break;
      case Encoding.DEFLATE_VALUE: {
        ByteString.Output output = ByteString.newOutput(request
            .getSerializedSize());
        DeflaterOutputStream delaterStream = new DeflaterOutputStream(output);
        request.writeTo(delaterStream);
        delaterStream.finish();
        delaterStream.close();
        encodedMessage.setDeflate(output.toByteString());
        break;
      }
      case Encoding.GZIP_VALUE: {
        ByteString.Output output = ByteString.newOutput(request
            .getSerializedSize());
        GZIPOutputStream gzipStream = new GZIPOutputStream(output);
        request.writeTo(gzipStream);
        gzipStream.finish();
        gzipStream.close();
        encodedMessage.setGzip(output.toByteString());
        break;
      }
      default:
        throw new IOException("Unsupported encoding: " + encoding);
    }
    return encodedMessage.build();
  }

  @Override
  public void send(Rpc rpc, RequestProto.Builder request) throws IOException {
    synchronized (this.lock) {
      this.registry.add(rpc.getReponseExtension());
      SocketRequestProto.Builder socketReq = SocketRequestProto.newBuilder();
      socketReq.setProtocolVersion(SocketProtocolVersion.V0100_VALUE);
      socketReq.setEncodedContent(pack(rpc.getTransportEncoding(),
          request.build()));
      socketReq.addAllAcceptEncoding(ACCEPTED_ENCODINGS);
      socketReq.setRequestId(this.callId);
      this.calls.put(this.callId, rpc);
      this.callId++;
      final SocketStreamProto reqFrame = SocketStreamProto.newBuilder()
          .addRequest(socketReq).build();
      reqFrame.writeTo(this.output);
      this.output.flush();
    }
  }

  private SocketResponseProto readOneResponse() throws IOException {
    synchronized (this.lock) {
      CodedInputStream codedInput = CodedInputStream.newInstance(this.input);
      final int tag = codedInput.readTag();
      if (tag == 0) {
        throw new IOException("Unexpected end of stream.");
      }
      final int type = tag & 0x07;
      if (type != WireFormat.WIRETYPE_LENGTH_DELIMITED) {
        codedInput.skipField(tag);
        throw new IOException("Wrong field type: " + type);
      }
      final int number = WireFormat.getTagFieldNumber(tag);
      if (number != SocketStreamProto.RESPONSE_FIELD_NUMBER) {
        SocketResponseProto.Builder socketResp = SocketResponseProto
            .newBuilder();
        socketResp.mergeFrom(codedInput, this.registry);
        return socketResp.build();
      } else {
        codedInput.skipField(tag);
        throw new IOException("Wrong field number: " + number);
      }
    }
  }

  private static ResponseProto getResponseOrThrow(SocketResponseProto socketResp)
      throws IOException {
    if (!socketResp.hasEncodedContent()) {
      throw new IOException("Missing encoded content.");
    }
    EncodedResponseProto encodedResponse = socketResp.getEncodedContent();
    ResponseProto response = null;
    if (encodedResponse.hasIdentity()) {
      response = encodedResponse.getIdentity();
    } else if (encodedResponse.hasDeflate()) {
      DeflaterInputStream deflaterStream = new DeflaterInputStream(
          encodedResponse.getGzip().newInput());
      response = ResponseProto.parseFrom(deflaterStream);
      deflaterStream.close();
    } else if (encodedResponse.hasGzip()) {
      GZIPInputStream gzipStream = new GZIPInputStream(encodedResponse
          .getGzip().newInput());
      response = ResponseProto.parseFrom(gzipStream);
      gzipStream.close();
    }
    if (response == null) {
      throw new IOException("Unknown content encoding.");
    }
    return response;
  }

  @Override
  public void run() {
    while (!this.end) {
      SocketResponseProto socketResp = null;
      Rpc rpc = null;
      synchronized (this.lock) {
        try {
          socketResp = readOneResponse();
          rpc = this.calls.get(new Long(socketResp.getRequestId()));
        } catch (IOException e) {
          e.printStackTrace();
        }
      }
      if (rpc != null) {
        try {
          rpc.setResponse(getResponseOrThrow(socketResp));
        } catch (IOException e) {
          rpc.setResponse(RpcStatus.RPC_ERROR, e.getMessage());
        }
      } else {
        System.err.println("No RPC for request: " + socketResp.getRequestId());
      }
    }
  }

  @Override
  public void close() throws IOException {
    synchronized (this.lock) {
      if (this.input != null) {
        try {
          this.input.close();
        } catch (IOException e) {}
        this.input = null;
      }
      if (this.output != null) {
        try {
          this.output.close();
        } catch (IOException e) {}
        this.output = null;
      }
      if (this.socket != null) {
        try {
          this.socket.close();
        } catch (IOException e) {}
        this.socket = null;
      }
    }
  }
}
